#include "ph.h"

#include <cuda.h>
#include "aesUtils.h"
#include "cudaUtils.h"
__constant__ byte DSBOX_INV[256];
__constant__ byte DSBOX[256];
__constant__ byte DSHIFT_NB[3][4][8];

//__device__ byte state[32];
//__device__ byte buffer[32];


void aesCipherInit()
{
    cuInit(0);
    CUDA_CHECK(cudaMemcpyToSymbol(DSBOX, SBOX, 256*sizeof(byte), 0, cudaMemcpyHostToDevice));
       
    CUDA_CHECK(cudaMemcpyToSymbol(DSBOX_INV, SBOX_INV, 256*sizeof(byte), 0, cudaMemcpyHostToDevice));
    
    CUDA_CHECK(cudaMemcpyToSymbol(DSHIFT_NB, SHIFT_NB, SHIFT_NB_SIZE, 0, cudaMemcpyHostToDevice));
    
}

__device__ void cudaAddRoundKey(byte* stateCell, byte* keyCell)
{
	*stateCell ^= *keyCell;
}

__device__ void cudaByteSub(byte* stateCell, int nullparam)
{
	*stateCell = DSBOX[*stateCell];
}

__device__ void cudaShiftRow(byte* stateCell, byte* buffer, int nb, int nullP)
{
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte dest_col = DSHIFT_NB[nb][row][col];
    byte* destCell = buffer + (dest_col<<2) + row;
    
    *destCell = *stateCell;
}

__device__ void cudaMixColumn(byte* state, byte* stateCell, byte* outState, byte* outStateCell, int nb)
{    
	const unsigned int& col = threadIdx.x;
    const unsigned int& row = threadIdx.y;
    
    byte* colCell = state + (col <<2);
    

    byte a0 = *(colCell);
    byte a1 = *(colCell+1);
    byte a2 = *(colCell+2);
    byte a3 = *(colCell+3);

    byte tmp, tm; 
    tmp = a0 ^ a1 ^ a2 ^ a3;
    
    byte* nextCell = colCell+ ( (row+1) & 0x3 );
    tm = *stateCell ^ *nextCell;
    tm = XTIME(tm);
    
    *outStateCell = *stateCell ^ tm ^ tmp;    
}

__device__ void cudaAesRound(byte* state, byte* buffer, int nb,
							byte* stateCell, byte* bufferCell, byte* keyCell)
{
    cudaByteSub(stateCell, 0);

	cudaShiftRow(stateCell, buffer, nb, 0);
	
    // now in cState is right State
    cudaMixColumn(buffer, bufferCell, state, stateCell, nb);

    //now in state is right state
    cudaAddRoundKey(stateCell, keyCell);
}

__device__ void cudaAesFinalRound(byte* state, byte* buffer, int nb,
							byte* stateCell, byte* bufferCell, byte* keyCell)
{
    cudaByteSub(stateCell, 0);

    cudaShiftRow(stateCell, buffer, nb, 0);
    
    cudaAddRoundKey(bufferCell, keyCell);
}

__device__ void cudaAESEncryptBlock(byte* state, byte* buffer, byte* roundKey, int nb, int nr)
{
	unsigned int wb = BLOCK_SIZE(nb)/4;
	const unsigned int& col = threadIdx.x;
    const unsigned int& row = threadIdx.y;
	
	byte* stateCell = state + (col<<2) + row;
	byte* bufferCell = buffer + (col<<2) + row;
	byte* keyCell = roundKey + (col<<2) + row;
    //--------------------------------------------------------
    
    cudaAddRoundKey(stateCell, keyCell);	
	
	
	for(int i = 1; i < nr; ++i)
    {
    	keyCell += (wb<<2);
    	cudaAesRound(state, buffer, nb, stateCell, bufferCell, keyCell);
    }
	
	keyCell += (wb<<2);
	
	cudaAesFinalRound(state, buffer, nb,
							stateCell, bufferCell, keyCell);
							
	//cipher is in bufferCell
	
}

#include "CudaAESOperations.cu"
#include "CudaAESInvOperations.cu"
#include "CudaAESFast.cu"
